[OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py - #2754

What does this PR do?

Type of change: New feature

hf_ptq.py had one distributed path, --use_fsdp2, which shards a dense model and replicates every expert of an MoE. This PR wires the expert-parallel path (the rest of the series; see the layout above) into it, so the documented CLI reaches an expert-parallel run: recipes, export formats, deprecations.

• Process group. main() runs inside distributed_run ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), which owns the process group on a distributed launch. The run takes this process's DistributedState (rank, size, device) from it, with --device for a single process. FSDP2 on one rank still shards through a process group, of one, which the load sets up ([OMNIML-6054] Route load_pretrained by the distribution plan #2779) and distributed_run destroys with the run. The example's own DistributedState, setup_distributed_args and process-group helpers are gone.
• One decision, read everywhere. load_model decides the plan once, with resolve_distribution_plan ([OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780, [OMNIML-6076] Describe expert parallelism in the model specs #2746), and loads every model by it through load_pretrained ([OMNIML-6054] Route load_pretrained by the distribution plan #2779): expert-parallel, FSDP2, get_model's whole-model loads and the low-memory load. load_pretrained records the plan, so the script records none itself: [OMNIML-6071] Calibrate a distributed model on per-rank data #2740's record after the --use_fsdp2 load is gone. Everything after reads the plan back from the model, not the flag or the process group: the calibration device, the data sharding, the calibration loop, and get_model's GPU placement.
  • MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE override the automatic choices.
  • --use_fsdp2 asks the decision for FSDP2 and no expert parallelism.
• Refusals. A plan that shards the whole model with FSDP2 raises without --use_fsdp2, the flag this script's FSDP2 checks hang on, which now run before anything is set up. resolve_distribution_plan refuses a plan that would quietly replicate an expert-parallel-capable MoE, and a model whose spec refuses expert parallelism (GPT-OSS, Llama-4, [OMNIML-6076] Describe expert parallelism in the model specs #2746) raises unless --use_fsdp2 rules expert parallelism out. Left to decide, the plan replicates a dense part it cannot FSDP2-shard within several expert-parallel groups, which the export cannot write yet; asking for that layout raises before loading and names the settings that avoid it. An AutoQuantize recipe raises on the expert-parallel and replicated routes, since its search scores each rank's shard without reducing the scores; --use_fsdp2 keeps main's warning. --low_memory_mode, --use_seq_device_map and --offload_folder (with it --max_cpu_memory_gb / --max_gpu_memory_gb) configure a single process's load, so an expert-parallel or replicated run refuses them before loading: an expert-parallel run never reaches those loaders, and a replicated one pins each rank's copy to its GPU, which they cannot work with. A single process and --use_fsdp2 keep main's handling. An expert-parallel run also refuses --vllm_fakequant_export before loading, since that exporter writes the whole model from each process it runs in; a replicated run keeps it (rank 0 exports), and --use_fsdp2 keeps its parsing-time refusal.
• Export. A model the plan replicates, such as a small dense model under torchrun, is exported by rank 0 alone; every rank used to write the same directory and race on config.json. A multimodal model's source config is rank 0's in every plan. A speculative-decoding model's export, which writes the whole checkpoint from one process, is refused on every rank when the model is sharded.
• Calibration.

  • Every distributed run calibrates on its own 1/world shard of the data ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740), the offline speculative-decoding loader included; a loader of a kind that cannot be split is refused.

  • Expert-parallel runs go through distributed_forward_loop, which re-plants the all-to-all forwards that quantizer insertion clobbers on transformers 5.16 and later. From transformers 5.16, quantizer insertion drops the routed experts' all-to-all forward wrapper, and a recipe whose calibration runs no loop (configs that need no calibration, such as int8_weight_only or w4a16_nvfp4, and the MX presets, whose algorithm is None) would leave it dropped, so the preview after PTQ would fail before the export. Once quantized, an expert-parallel run therefore re-plants it. On 5.15 and earlier the TP transforms are forward hooks that insertion leaves in place, and the re-plant finds nothing to restore.

  • An expert-parallel VLM calibrated on text runs its batches through the language model, but the loop re-plants the all-to-all forwards over the whole model: transformers keeps the expert-parallel plan on the model from_pretrained built, and from 5.16 a VLM's language model reports none of its own, so a loop handed the language model would re-plant nothing, and the first batch would raise for the missing EP process group. The --qformat path of an expert-parallel VLM also needs [OMNIML-6071] Calibrate a distributed model on per-rank data #2740's guard fix, which lets a module with no routed experts (the vision tower the script quantizes with every quantizer off) pass the expert-parallel calibration check.

  • FSDP2 keeps main's plain loop, with its split-on-OOM retry.

  • The calibration algorithms under a plan that spans several ranks:

    	Algorithms
    Supported and run on GPUs	max (amax all-reduced over the data-, expert- and tensor-parallel groups), mse (searches the weights max leaves identical on every rank)
    Supported, not yet run on GPUs	gptq (Hessians merged across each module's data-parallel replicas), smoothquant (scales from max's synced amax), nvfp4_act_headroom (activation scales MAX-reduced over data parallelism)
    Refused before quantizing ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740)	awq_lite, awq_clip, awq_full, svdquant, local_hessian, imatrix, and max or mse with distributed_sync off, also as another algorithm's weight or scale algorithm: each keeps its statistics on its rank, so the ranks' copies of a weight would come out different
    Not assessed	lsq's learnable scales (its first scale calibration follows the rows above)

• Previews keep a sharded model's ranks together. Each rank previews the first sample of its own calibration shard, so the ranks finish decoding at different steps. Under expert parallelism or FSDP2 every forward runs a collective, and a rank that stopped decoding left the others blocked in it: Qwen3-30B-A3B at NVFP4 and EP4 hung in the post-PTQ preview. The previews, Nemotron-VL's included, now generate with synced_gpus when the model's plan shards it. Under expert parallelism a preview that fails raises, which aborts the run, rather than leaving the other ranks in the next step's collective; its error says the export was not reached and names --skip_generate. --use_fsdp2 keeps main's best-effort skip and still exports, which is safe for a failure every rank meets at the same step. A failure on one rank alone leaves the others in generate's collectives until the collective timeout. While a preview decodes, generate is kept from switching a MoE's experts to batched_mm, which bypasses the quantizers, on every module of the model that has the decode hook: a model can generate through a submodule, as a Nemotron-VL model's chat generates through its language model. The guard puts back the hook each module held itself, so a hook bound per instance survives and previews nest. The ranks also decide together whether to run Nemotron-VL's image preview: each node reads the sample images from its own checkout, and the repository ships none.
• GPU placement. When the plan gives every one of several ranks the whole model, each rank pins it to its own GPU, on every load path (the pack-quantized and MXFP4 loads, DiffusionGemma's sequential map), and --use_seq_device_map, which spreads one model over every GPU, and --offload_folder are refused before loading there. Under a plan with several ranks, a device that names no GPU index is refused before loading, rather than left to accelerate's "auto". A pinned copy is not sized against the node's GPUs, a limit its one-device map would ignore. A single process keeps spreading a model too big for one GPU, as before.
• Docs. README section "Expert-Parallel Post-Training Quantization for MoE Models": how to launch, what the knobs override, which models are refused, and that GPT-OSS and Llama-4 (refused by expert parallelism, and by --use_fsdp2 for their released MXFP4 / multimodal checkpoints) quantize in a single process.
• Memory report fix. launch_memory_monitor probed NVML before initializing it. Every probe raised Uninitialized, so the per-GPU peak-memory report never started on any host where nothing else had initialized NVML. The probe now initializes NVML itself, and only rank 0 starts the monitor, which reads every GPU on its node.

Usage

# One node, 8 GPUs: the experts shard over all 8 ranks, each calibrating its own data.
torchrun --nproc_per_node 8 examples/hf_ptq/hf_ptq.py \
    --pyt_ckpt_path Qwen/Qwen3-30B-A3B --qformat nvfp4 --export_path /out/qwen3-30b-a3b-nvfp4

Testing

• New tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py (70 cases; the decision itself is tested in [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 and [OMNIML-6076] Describe expert parallelism in the model specs #2746):
  • what the script adds to the decision: FSDP2 without expert parallelism for --use_fsdp2, and FSDP2 sharding refused unless --use_fsdp2;
  • load_model loading by the plan: an EP plan and --use_fsdp2 through load_pretrained and no other loader, a data-parallel one through get_model;
  • per-rank GPU pinning by the plan, and single-process placement left alone;
  • the single-process load options (--low_memory_mode, --use_seq_device_map, --offload_folder) refused on the expert-parallel and replicated routes and kept by a single process;
  • main() running inside distributed_run, which it leaves even when the run raises, and handing the run the state distributed_run yields, with --device for a single process;
  • per-rank calibration data;
  • the distributed loop, with the EP mesh, for an expert-parallel plan, and the plain one for FSDP2.
  • the previews, before and after quantization and Nemotron-VL's, generating with synced_gpus exactly when the plan shards the model (EP or FSDP2 over several ranks);
  • the image preview skipped on every rank when one has no sample image;
  • exporting a replicated model from rank 0 only, and every rank of a sharded one;
  • a sharded speculative-decoding model refused at export on every rank, with nothing written;
  • the plan the export cannot write refused before loading, while the world-size degree still resolves;
  • the vLLM fake-quant export refused on the expert-parallel route and kept by replicated and single-process runs;
  • the all-to-all forwards re-planted after a quantization that runs no calibration loop, for an expert-parallel plan only. A 2-GPU A/B on an earlier revision of this branch (tiny Qwen3-MoE, EP 2 with the dense part replicated, MXFP8 on weight and input quantizers with algorithm None, previews on): on transformers 5.18 the run without the re-plant failed in the post-quantization preview on both ranks ("torch_a2a experts forward requires the EP process_group") and wrote no checkpoint, and with it the run completed, restoring 8 forward wrappers, and exported; on transformers 5.9.0, where there is nothing to re-plant (up to 5.15 the TP transforms are hooks that quantizer insertion leaves alone), both runs exported byte-identical files;
  • a tiny Qwen3-VL-MoE, as the expert-parallel load leaves it, calibrated on text over a one-rank gloo group (fails without the fix on transformers 5.18; skipped below 5.16). On GPUs at this revision, a tiny Kimi-K3, a VLM MoE (bf16, 32 experts), on 2 GPUs at EP 2 with transformers 5.18, NVFP4 with an FP8 KV cache, calibrated on text and previewed before and after PTQ on both ranks: the calibration loop re-installed its 3 experts' forward wrappers, and the run exported;
  • Mutation-checked: each test fails under a deliberate break of the code it covers.
• New 2-GPU smoke test test_ptq_qwen3_moe_expert_parallel (tests/examples/hf_ptq/test_llm_ptq.py): torchrun --nproc_per_node 2 hf_ptq.py with MODELOPT_EP_DEGREE=2 on a tiny Qwen3-MoE and the FP8 recipe. It checks that every expert of both ranks is exported FP8 with its weight scale, and that the config names FP8. It is the tests' one real expert-parallel run of the script; the file above covers the wiring with stand-ins. It passed on 2 GB200s in 29.5 s at f4bb395488 (2026-10-16), which carries this test on a revision before the latest loader and calibration fixes ([OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6071] Calibrate a distributed model on per-rank data #2740). It skips on CI's 1-GPU PR runner and runs on 2-GPU runners.
• tests/examples/hf_ptq/test_hf_ptq_model_type.py (main's): the preview tests' stand-in models are modules now, since the previews read their distribution plan.
• tests/examples/hf_ptq/test_example_utils.py: the preview keeping generate from switching the experts implementation, also on a language model that a model around it generates through (fails without the fix), a model without the decode hook previewed as it is, and a decode hook the model holds itself put back, also when previews nest.
• New tests/unit/torch/utils/test_memory_monitor.py, against a fake reference-counted NVML: the monitor starts when NVML can report memory (this fails on the old code), and is skipped with NVML shut down when it cannot. tests/examples/hf_ptq/test_hf_ptq_args.py: the memory monitor started on rank 0 only.
• Plan tests the library's and get_model's tests cover are dropped, and the load-call, calibration-loop and memory-monitor pairs are one parametrized test each. Each drop is checked by mutation.
• test_hf_ptq_expert_parallel.py also covers: the pin on the pack-quantized, MXFP4 and DiffusionGemma load paths and the sequential-map refusal; a pinned copy not held to a node-wide memory limit; the offline speculative-decoding loader split across ranks; a failed preview raising on an expert-parallel model and skipped, with the export still run, on a replicated or --use_fsdp2 one; a device that names no GPU index refused where each rank holds its own copy; and an AutoQuantize recipe refused on the routes this script adds.
• The DeepSeek-V3 multi-rank cases of test_hf_ptq_expert_parallel.py skip below transformers 5.7, where [OMNIML-6076] Describe expert parallelism in the model specs #2746 refuses expert parallelism before the route they exercise; their single-process and --use_fsdp2 cases still run there.
• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.
• The GPU results below were measured on earlier revisions of this branch, before this script's loads moved onto load_pretrained ([OMNIML-6054] Route load_pretrained by the distribution plan #2779, [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780); this revision's GPU tests run in CI only. The series' GPU tests on this branch: 19 passed on a 4-GPU node. On 1 GPU, CI's PR runner, 6 passed and 13 skipped. They cover:
  • the expert-parallel load with the dense part replicated and sharded, at EP 4 and EP 2;
  • amax sync;
  • the distributed export.
• Against the series' base, on one GPU node:
  • nox -s gpu (all of tests/gpu, 1 GPU, CI's PR-runner shape): every test failing on this branch also fails at the base, and the only new tests are the series' own.
  • The hf_ptq example lane (tests/examples/hf_ptq, TensorRT-LLM image, 1 GPU): no outcome changes on the 296 shared tests. The series' 30 pass, except 2 that need transformers 5.7.
  • tests/gpu_megatron (2 GPUs): no outcome changes over the 311 tests both runs reached. The rest time out identically at the base in that image.
• Expert-parallel runs end to end. All of them predate [OMNIML-6073] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's expert weight-scale fix: on every rank but rank 0, an expert's weight scale came out too large wherever rank 0's expert at the same index had the larger one, which costs precision but never clips; the fix can only tighten such scales.

  • Qwen3-30B-A3B at EP16 over 4 nodes, NVFP4 and FP8: the exports with the dense part FSDP-sharded over the world and with it replicated have byte-identical weights and weight scales; only calibration-derived input scales differ. Every dense FP8 weight is FP8 with its scale, and no NVFP4 weight has the all-zero high nibbles of a lost pack shift.

  • Qwen3-VL-30B and MiniMax-M3: tensor counts match single-GPU exports, and every expert weight is packed.

  • MiniMax-M2.7 (block-FP8 source) at EP16 over 4 nodes, NVFP4: every exported weight name is the source checkpoint's, and routed experts are block_sparse_moe.experts.N.w1/w2/w3. Served by vLLM (0.30 dev build, TP4), it scores 0.94 / 0.93 on gsm8k (5-shot, 100 samples, flexible / strict), against the FP8 source's 0.93 / 0.95.

  • DeepSeek-V3 and GLM: on an earlier revision.

  • Kimi-K3 NVFP4 at EP64 over 16 GB200 nodes, in 10 minutes (peak 155 GiB per GPU on an earlier build of this series), checked against the source checkpoint:

    • the router correction biases match;
    • routed and shared experts carry the names vLLM 0.28's loader reads;
    • dense weights dequantize back to the BF16 source at cosine 0.9955;
    • the remote code, tokenizer and processor files ship with the export;
    • the modules vLLM builds unquantized stay BF16, as in nvidia/Kimi-K3-NVFP4.

    The same checks flag the export made before this series' fixes. vLLM 0.28 loads the export and starts serving, but it cannot serve a Kimi-K3 NVFP4 export with an FP8 KV cache: its FP8-KV path needs an FP8 prefill query, which its NVFP4 linears reject. TBD: generation and gsm8k from the checkpoint-mirror recipe models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention on this series' code with the [OMNIML-6073] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 fix.

• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.
• Review fixes, 2026-10-27: the Whisper calibration loader stays whole on every rank, so the preview matches its printed transcript. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.
• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). hf_ptq keeps main's removal of the FSDP2 KV-cache AutoQuantize refusal ([OMNIML-5987] [2/3] Support DP/FSDP2 KV AutoQuant and checkpoint resume #2664); main's new test that such a run reaches the distributed load now stubs load_pretrained, which this PR loads every distributed plan through, and an import that only the removed refusal test used is gone.
• Review fixes, 2026-10-27 (Claude's review after the rebase): --use_fsdp2's checks (validate_fsdp2_supported) run on the checkpoint's config before the plan is decided, as they ran before this PR, so a checkpoint they refuse (VILA, multimodal, pack-quantized) gets their message instead of the meta-device build's error. The new test fails without the change.
• Review suggestions, 2026-10-27 evening: hf_ptq refuses a layerwise-export recipe on a multi-rank launch, before the load; the memory monitor runs on each node's first rank; the unreachable Whisper loader exemption is gone; the off-GPU warning covers a replicated rank's device.
• Review suggestions, 2026-10-27 night: one helper, layerwise_export_requested, decides both the pre-load refusal and the layerwise export switch; a comment names the load paths without a device map (Nemotron VL), where each rank builds a whole host copy.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Single-process and --use_fsdp2 runs behave as before, except that under torchrun a --device other than cuda is now refused (by distributed_run, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), where main's --use_fsdp2 ignored it and ran each rank on cuda:<local_rank>. A torchrun launch without --use_fsdp2 now runs an MoE expert-parallel.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ New feature, 0.49.0.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 21 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2749.

Summary by CodeRabbit

• New Features
  • Added distributed, expert-parallel post-training quantization for Hugging Face MoE models, with calibration data distributed across ranks and checkpoint shards written in parallel.
  • Added options to control expert-parallel degree and dense-layer FSDP2 sharding, with validation for unsupported configurations.
  • Added synchronized preview generation for distributed runs.
• Bug Fixes
  • Improved GPU memory monitoring when NVML initialization or memory reporting is unavailable.
• Documentation
  • Added setup guidance, launch examples, requirements, and distribution options for expert-parallel quantization.